# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import datetime
import pathlib
import platform
import numpy as np
import pandas as pd

import dateutil.parser
from absl import app, flags

from coin.base.datetime_util import iterate_date
from coin.experimental.xguo.convert_symbol_to_product import convert_symbol_to_product
from coin.experimental.xguo.fruit.util.feed_util import (AutoInc,
                                                         load_products_feed_as_json,
                                                         get_moving_average,
                                                         get_moving_diff,
                                                         get_moving_return,
                                                         get_mid_series,
                                                         get_trade_qty_series,
                                                         get_moving_sum,
                                                         get_moving_abs_sum,
                                                         get_depth,
                                                         remove_nan_entries)

#matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


class Category(enum.Enum):
  LONG = 'long'
  SHORT = 'short'
  HOLD_POSITION = 'hold_position'
  CLOSE_POSITION = 'close_position'


def plot_result(
    mid_price_series,
    trade_qty_series,
    buy_depth,
    sell_depth,
):
  plt.figure(figsize=(20, 10))
  ser = get_moving_sum(trade_qty_series, 300)

  plt.subplot(4, 1, 1)
  mid_price_series.plot()

  plt.subplot(4, 1, 2)
  freq = datetime.timedelta(seconds=300)
  ser = ser.asfreq(freq)
  ser = remove_nan_entries(ser)
  print(ser)
  fft_ser = np.fft.fft(ser)
  fft_ser = np.fft.fftshift(np.abs(fft_ser))
  plt.plot(fft_ser)

  plt.subplot(4, 1, 3)
  plt.stem(ser.index, ser)

  plt.subplot(4, 1, 4)
  buy_depth.plot()
  sell_depth.plot()
  plt.show()


def gen_x_value(mid_price, trade_qty_series):
  record = {}
  for win in [1, 5, 10, 30, 60, 100, 200, 300, 600, 1200, 2400]:
    window = datetime.timedelta(seconds=win)

    label = f's{win}_price'
    record[label] = get_moving_average(mid_price, window)

    label = f's{win}_sum'
    record[label] = get_moving_sum(trade_qty_series, window)

    label = f's{win}_abs_sum'
    record[label] = get_moving_abs_sum(trade_qty_series, window)
  return pd.DataFrame(record)


def main(_):
  if FLAGS.home_dir is None:
    if 'hive' in platform.node():
      home_dir = pathlib.Path('/remote/iosg/home/xguo')
    else:
      home_dir = pathlib.Path.home()
  else:
    home_dir = pathlib.Path(FLAGS.home_dir)

  root_dir = home_dir.joinpath('converted_book')

  products = []
  for symbol in FLAGS.symbol.split(','):
    product = convert_symbol_to_product(symbol)
    products.append(product)

  if '-' in FLAGS.trading_date:
    start_date, end_date = FLAGS.trading_date.split('-')
    start_date = dateutil.parser.parse(start_date)
    end_date = dateutil.parser.parse(end_date)
    trading_dates = list(iterate_date(start_date, end_date))
  else:
    trading_dates = [dateutil.parser.parse(FLAGS.trading_date)]

  feeds = load_products_feed_as_json(root_dir, trading_dates, products)
  prices = {}
  trade_qty = {}
  for product, feed_list in feeds.items():
    price = get_mid_series(feed_list)
    trade_qty = get_trade_qty_series(feed_list)
    buy_depth, sell_depth = get_depth(feed_list, num_depth=1)
    plot_result(price, trade_qty, buy_depth, sell_depth)
    # df = gen_x_value(price, trade_qty)
    # df.to_pickle('')
    # print(df.to_string(line_width=1000))


if __name__ == '__main__':
  flags.DEFINE_string(
      'symbol',
      None,
      'Specify product full symbol.',
  )

  flags.DEFINE_string('trading_date', None, 'Trading date.')

  flags.DEFINE_string('home_dir', None, 'Home directory.')
  app.run(main)
